import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
import numpy as np
import os
import time

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

def check_device():
    """检查并返回可用的设备"""
    print("=== 设备检查 ===")
    # 检查MPS支持，同时处理兼容性问题
    if (hasattr(torch.backends, 'mps') and 
        torch.backends.mps.is_available() and 
        torch.backends.mps.is_built()):
        try:
            # 尝试创建一个MPS张量来测试是否真正可用
            test_tensor = torch.zeros(1, device='mps')
            device = torch.device("mps")
            print("使用MPS设备 (Apple Silicon GPU)")
        except RuntimeError as e:
            print(f"MPS测试失败: {e}")
            device = torch.device("cpu")
            print("回退到CPU设备")
    elif torch.cuda.is_available():
        device = torch.device("cuda")
        print("使用CUDA设备 (NVIDIA GPU)")
    else:
        device = torch.device("cpu")
        print("使用CPU设备")
    print(f"设备: {device}")
    return device

def load_data():
    """加载并预处理MNIST数据"""
    print("\n=== 加载MNIST数据 ===")
    
    # 数据预处理
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))  # MNIST数据集的均值和标准差
    ])
    
    # 下载并加载训练数据
    trainset = torchvision.datasets.MNIST(
        root='./mnist_data', 
        train=True, 
        download=True, 
        transform=transform
    )
    trainloader = torch.utils.data.DataLoader(
        trainset, 
        batch_size=32,  # 进一步减小批次大小
        shuffle=True,
        num_workers=0,  # 设置为0避免多进程问题
        pin_memory=False  # 关闭pin_memory以避免MPS问题
    )
    
    # 下载并加载测试数据
    testset = torchvision.datasets.MNIST(
        root='./mnist_data', 
        train=False, 
        download=True, 
        transform=transform
    )
    testloader = torch.utils.data.DataLoader(
        testset, 
        batch_size=32,  # 进一步减小批次大小
        shuffle=False,
        num_workers=0,  # 设置为0避免多进程问题
        pin_memory=False  # 关闭pin_memory以避免MPS问题
    )
    
    print(f"训练集大小: {len(trainset)}")
    print(f"测试集大小: {len(testset)}")
    return trainloader, testloader

class CNN(nn.Module):
    """卷积神经网络模型"""
    def __init__(self):
        super(CNN, self).__init__()
        self.conv1 = nn.Conv2d(1, 16, kernel_size=3, padding=1)  # 减少通道数
        self.conv2 = nn.Conv2d(16, 32, kernel_size=3, padding=1)  # 减少通道数
        self.pool = nn.MaxPool2d(2, 2)
        self.fc1 = nn.Linear(32 * 7 * 7, 64)  # 减少神经元数量
        self.fc2 = nn.Linear(64, 10)
        self.relu = nn.ReLU()
        self.dropout = nn.Dropout(0.1)  # 减少dropout比例
        
    def forward(self, x):
        x = self.pool(self.relu(self.conv1(x)))  # (batch, 16, 14, 14)
        x = self.pool(self.relu(self.conv2(x)))  # (batch, 32, 7, 7)
        x = x.view(-1, 32 * 7 * 7)               # (batch, 32*7*7)
        x = self.relu(self.fc1(x))               # (batch, 64)
        x = self.dropout(x)
        x = self.fc2(x)                          # (batch, 10)
        return x

def train_model(model, trainloader, testloader, device, epochs=2):  # 进一步减少训练轮数
    """训练模型"""
    print(f"\n=== 开始训练 (设备: {device}) ===")
    
    # 损失函数和优化器
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=0.001)
    
    # 将模型移到指定设备
    try:
        model.to(device)
    except Exception as e:
        print(f"模型迁移到设备失败: {e}")
        print("尝试使用CPU...")
        device = torch.device("cpu")
        model.to(device)
    
    # 训练历史记录
    train_losses = []
    train_accuracies = []
    test_accuracies = []
    
    # 训练过程
    for epoch in range(epochs):
        model.train()
        running_loss = 0.0
        correct = 0
        total = 0
        
        start_time = time.time()
        for i, (inputs, labels) in enumerate(trainloader):
            try:
                # 将数据移到指定设备
                inputs, labels = inputs.to(device), labels.to(device)
                
                # 对于MPS设备，需要特别处理以避免段错误
                if device.type == "mps":
                    # 在MPS上需要确保数据在正确的设备上且连续
                    inputs = inputs.contiguous()
                    labels = labels.contiguous()
                
                # 梯度清零
                optimizer.zero_grad()
                
                # 前向传播
                outputs = model(inputs)
                loss = criterion(outputs, labels)
                
                # 反向传播
                loss.backward()
                optimizer.step()
                
                # 统计信息
                running_loss += loss.item()
                _, predicted = torch.max(outputs.data, 1)
                total += labels.size(0)
                correct += (predicted == labels).sum().item()
                
                # 打印进度（进一步减少打印频率）
                if i % 300 == 299:
                    print(f'Epoch [{epoch+1}/{epochs}], Step [{i+1}/{len(trainloader)}], '
                          f'Loss: {running_loss/300:.4f}')
                    running_loss = 0.0
                    
            except RuntimeError as e:
                if "MPS" in str(e):
                    print(f"MPS运行时错误，跳过当前批次: {e}")
                    # 清理缓存
                    if hasattr(torch.mps, 'empty_cache'):
                        torch.mps.empty_cache()
                    continue
                else:
                    raise e
        
        epoch_time = time.time() - start_time
        
        # 计算训练准确率
        train_acc = 100 * correct / total
        train_losses.append(loss.item() if 'loss' in locals() else 0)
        train_accuracies.append(train_acc)
        
        # 计算测试准确率
        test_acc = evaluate_model(model, testloader, device)
        test_accuracies.append(test_acc)
        
        print(f'Epoch [{epoch+1}/{epochs}] 完成 - '
              f'耗时: {epoch_time:.2f}s, '
              f'训练准确率: {train_acc:.2f}%, '
              f'测试准确率: {test_acc:.2f}%')
        
        # 对于MPS设备，添加同步和清理操作以避免段错误
        if device.type == "mps":
            if hasattr(torch.mps, 'synchronize'):
                try:
                    torch.mps.synchronize()
                except:
                    pass
            if hasattr(torch.mps, 'empty_cache'):
                try:
                    torch.mps.empty_cache()
                except:
                    pass
    
    return train_losses, train_accuracies, test_accuracies

def evaluate_model(model, dataloader, device):
    """评估模型"""
    model.eval()
    correct = 0
    total = 0
    
    with torch.no_grad():
        for inputs, labels in dataloader:
            try:
                # 将数据移到指定设备
                inputs, labels = inputs.to(device), labels.to(device)
                
                # 对于MPS设备，需要特别处理
                if device.type == "mps":
                    inputs = inputs.contiguous()
                    labels = labels.contiguous()
                
                outputs = model(inputs)
                _, predicted = torch.max(outputs.data, 1)
                total += labels.size(0)
                correct += (predicted == labels).sum().item()
                
                # 对于MPS设备，添加同步操作
                if device.type == "mps":
                    if hasattr(torch.mps, 'synchronize'):
                        try:
                            torch.mps.synchronize()
                        except:
                            pass
                        
            except RuntimeError as e:
                if "MPS" in str(e):
                    print(f"评估时MPS运行时错误: {e}")
                    continue
                else:
                    raise e
    
    accuracy = 100 * correct / total
    return accuracy

def visualize_samples(trainloader, num_samples=9):
    """可视化数据样本"""
    print("\n=== 可视化数据样本 ===")
    
    # 获取一批数据
    dataiter = iter(trainloader)
    images, labels = next(dataiter)
    
    # 可视化图像
    fig = plt.figure(figsize=(8, 8))
    for i in range(num_samples):
        plt.subplot(3, 3, i+1)
        # 反归一化图像
        img = images[i] * 0.3081 + 0.1307
        plt.imshow(img.numpy().squeeze(), cmap='gray')
        plt.title(f'标签: {labels[i].item()}')
        plt.axis('off')
    plt.tight_layout()
    plt.show()

def predict_samples(model, testloader, device, num_samples=5):  # 减少预测样本数
    """预测示例"""
    print("\n=== 预测示例 ===")
    
    model.eval()
    dataiter = iter(testloader)
    images, labels = next(dataiter)
    
    try:
        # 将数据移到指定设备
        images, labels = images.to(device), labels.to(device)
        
        # 对于MPS设备，需要特别处理
        if device.type == "mps":
            images = images.contiguous()
            labels = labels.contiguous()
        
        # 预测
        with torch.no_grad():
            outputs = model(images[:num_samples])
            _, predicted = torch.max(outputs, 1)
        
        # 打印预测结果
        print("前5个测试样本的预测结果:")
        for i in range(num_samples):
            print(f"样本 {i+1}: 真实标签={labels[i].item()}, 预测标签={predicted[i].item()}")
            
    except RuntimeError as e:
        if "MPS" in str(e):
            print(f"预测时MPS运行时错误: {e}")
        else:
            raise e

def main():
    """主函数"""
    print("MNIST手写数字识别 - PyTorch卷积神经网络实现 (支持MPS)")
    print("=" * 60)
    
    # 检查设备
    device = check_device()
    
    # 加载数据
    trainloader, testloader = load_data()
    
    # 可视化样本数据
    try:
        visualize_samples(trainloader)
    except Exception as e:
        print(f"可视化数据时出错: {e}")
        print("跳过可视化步骤...")
    
    # 创建模型
    print("\n=== 创建卷积神经网络模型 ===")
    model = CNN()
    print(f"模型参数数量: {sum(p.numel() for p in model.parameters())}")
    
    # 训练模型
    try:
        train_losses, train_accuracies, test_accuracies = train_model(
            model, trainloader, testloader, device, epochs=2
        )
    except Exception as e:
        print(f"训练过程中出现错误: {e}")
        print("程序将继续执行后续步骤...")
    
    # 预测示例
    try:
        predict_samples(model, testloader, device)
    except Exception as e:
        print(f"预测示例时出错: {e}")
    
    # 保存模型
    try:
        torch.save(model.state_dict(), '../pytorch_cnn_mnist_model.pth')
        print("\n模型已保存为 'pytorch_cnn_mnist_model.pth'")
    except Exception as e:
        print(f"保存模型时出错: {e}")
        print("模型未保存")
    
    print("\n" + "=" * 60)
    print("程序执行完成!")

if __name__ == "__main__":
    main()